#include "StdAfx.h"
#include "AVISplitter.h"
#include "..\WMP\Common.h"

AVISplitterFactory::AVISplitterFactory(IClassFactory *factory) {
	this->factory = factory;
}
HRESULT AVISplitterFactory::Create(IClassFactory *factory, REFIID riid, void** ppv) {
	if (!ppv || !factory)
		return E_POINTER;
	AVISplitterFactory *wrap = new AVISplitterFactory(factory);
	if (!wrap)
		return E_OUTOFMEMORY;
	//if (IsEqualIID(riid, IID_IClassFactory))
	*ppv = (IClassFactory*)wrap;
	return S_OK;
}
AVISplitterFactory::~AVISplitterFactory() {
}

HRESULT STDMETHODCALLTYPE AVISplitterFactory::QueryInterface(REFIID riid, void **ppvObject) {
	if(!ppvObject) return E_POINTER;
	if (IsEqualIID(riid, IID_IClassFactory))
		*ppvObject = (IClassFactory*)this;
	else
		return this->factory->QueryInterface(riid, ppvObject);
	this->factory->AddRef();
	return NOERROR;
}

ULONG STDMETHODCALLTYPE AVISplitterFactory::AddRef() {
	return this->factory->AddRef();
}

ULONG STDMETHODCALLTYPE AVISplitterFactory::Release() {
	ULONG ref = this->factory->Release();
	if (!ref) {
		delete this;
		return ULONG(0);
	} else
		return ref;
}

HRESULT STDMETHODCALLTYPE AVISplitterFactory::CreateInstance(IUnknown *pUnkOuter, REFIID riid, void **ppvObject) {
	CreateTextBuffer(textBuffer);
	CreatePathBuffer(idBuffer);

	if (IsEqualIID(riid, IID_IUnknown)) {
		return AVISplitter::Create(this->factory, ppvObject);
	} else {
		_tcscpy(textBuffer, _T("WARNING: Unknown AviSplitter interface query "));
		LPOLESTR id;
		if (SUCCEEDED(StringFromIID(riid, &id))) {
			_tcscat(textBuffer, FromLPWSTR(id, idBuffer, idBufferLength));
			CoTaskMemFree(id);
		}
		OutputDebugString(textBuffer);
	}
	return this->factory->CreateInstance(pUnkOuter, riid, ppvObject);
}

HRESULT STDMETHODCALLTYPE AVISplitterFactory::LockServer(BOOL fLock) {
	return this->factory->LockServer(fLock);
}

AVISplitter::AVISplitter()
	: streams(NULL), inner(NULL), splitter(NULL), graph(NULL), ref(1) { }
HRESULT AVISplitter::Create(IClassFactory *factory, void** ppvObject) {
	if (!factory)
		return E_POINTER;
	AVISplitter *wrap = new AVISplitter();
	IUnknown *outer = (IUnknown*)wrap;
	HRESULT res = factory->CreateInstance(outer, IID_IUnknown, (void**)&wrap->inner);
	if (SUCCEEDED(res)) {
		if (FAILED(res = wrap->inner->QueryInterface(IID_IBaseFilter, (void**)&wrap->splitter)))
			wrap->inner->Release();
		else
			outer->Release(); //last query interface increments ref counter
	}
	*ppvObject = outer;
	return res;
}
AVISplitter::~AVISplitter() {
	delete this->streams;
	this->ref = ULONG(1);
	IUnknown* outer = this;
	outer->AddRef();
	FreeObject(this->splitter);
	FreeObject(this->inner);
	FreeObject(this->graph);
}

HRESULT STDMETHODCALLTYPE AVISplitter::QueryInterface(REFIID riid, void **ppvObject) {
	if(!ppvObject) return E_POINTER;
	if (IsEqualIID(riid, IID_IAMStreamSelect))
		*ppvObject = (IAMStreamSelect*)this;
	else if (IsEqualIID(riid, IID_IUnknown))
		*ppvObject = (IUnknown*)this;
	else
		return this->inner->QueryInterface(riid, ppvObject);
	AddRef();
	return NOERROR;
}

ULONG STDMETHODCALLTYPE AVISplitter::AddRef() {
	return InterlockedIncrement(&this->ref);
}

ULONG STDMETHODCALLTYPE AVISplitter::Release() {
	if (!this->ref)
		return ULONG(0);
	ULONG ref = InterlockedDecrement(&this->ref);
	if (!ref) {
		delete this;
		return ULONG(0);
	} else
		return ref;
}

void FreePinMediaType(AM_MEDIA_TYPE& mt) {
	if (mt.cbFormat != 0) {
		CoTaskMemFree(mt.pbFormat);
		mt.cbFormat = 0;
		mt.pbFormat = NULL;
	}
	// Unecessary because pUnk should not be used, but safest.
	FreeObject(mt.pUnk);
}
HRESULT CopyPinMediaType(AM_MEDIA_TYPE *pmtTarget, const AM_MEDIA_TYPE *pmtSource) {
	*pmtTarget = *pmtSource;
	if (pmtSource->cbFormat != 0 && pmtSource->pbFormat) {
		pmtTarget->pbFormat = (PBYTE)CoTaskMemAlloc(pmtSource->cbFormat);
		if (!pmtTarget->pbFormat) {
			pmtTarget->cbFormat = 0;
			return E_OUTOFMEMORY;
		}
		CopyMemory(pmtTarget->pbFormat, pmtSource->pbFormat, pmtTarget->cbFormat);
	}
	if (pmtTarget->pUnk) {
		pmtTarget->pUnk->AddRef();
	}
	return S_OK;
}

AVISplitter::PinInfo* AVISplitter::PinInfo::CreateNext() {
	IEnumPins* pins;
	IPin* pin2;
	PIN_DIRECTION dir;
	AVISplitter::PinInfo* info = NULL;

	if (!this->filter || FAILED(this->filter->EnumPins(&pins))) return NULL;
	try {
		while(pins->Next(1u, &pin2, NULL) == S_OK) {
			if (FAILED(pin2->QueryDirection(&dir)) || dir != PINDIR_OUTPUT) {
				pin2->Release();
				continue;
			}
			AM_MEDIA_TYPE *pMediaType = (AM_MEDIA_TYPE *)CoTaskMemAlloc(sizeof(AM_MEDIA_TYPE));
			if (pMediaType != NULL) {
				if (FAILED(pin2->ConnectionMediaType(pMediaType))) {
					CoTaskMemFree(pMediaType);
					pMediaType = NULL;
				}
			}
			info = new AVISplitter::PinInfo(this->index, pin2, pMediaType, false);
			break;
		}
	} finally (
		FreeObject(pins);
	)
	return info;
}

AVISplitter::PinInfo::PinInfo(long index, IPin* pin, AM_MEDIA_TYPE* type, bool chain)
	: connect(NULL), filter(NULL), next(NULL) {
	PIN_INFO info;
	FILTER_INFO info2;

	this->index = index;
	this->pin = pin;
	this->type = type;
	if (this->pin)
		this->pin->ConnectedTo(&this->connect);
	if (this->connect && chain && SUCCEEDED(this->connect->QueryPinInfo(&info))) {
		this->filter = info.pFilter;
		if (this->filter) {
			if (SUCCEEDED(this->filter->QueryFilterInfo(&info2))) {
				wcsncpy(this->filterName, info2.achName, MAX_FILTER_NAME);
				info2.pGraph->Release();
			} else this->filterName[0] = 0;
			this->next = CreateNext();
		}
	}
}

AVISplitter::PinInfo::~PinInfo() {
	FreeObject(this->pin);
	FreeObject(this->connect);
	FreeObject(this->filter);
	if (this->type) {
		FreePinMediaType(*this->type);
		CoTaskMemFree(this->type);
	}
}

HRESULT AVISplitter::EnsureGraph() {
	FILTER_INFO info;

	if (this->graph) return S_OK;
	HRESULT res = this->splitter->QueryFilterInfo(&info);
	if (SUCCEEDED(res)) {
		this->graph = info.pGraph;
		if (!this->graph)
			res = E_POINTER;
	}
	return res;
}

HRESULT AVISplitter::EnsurePins() {
	HRESULT res;
	IEnumPins* pins = NULL;
	IPin* pin = NULL;
	PIN_DIRECTION dir;
	CreateTextBuffer(textBuffer);

	if (this->streams)
		return S_OK;
	this->streams = new CArray<PinInfo*>();
	if (!this->streams)
		return E_OUTOFMEMORY;
	if (FAILED(res = this->splitter->EnumPins(&pins)))
		return res;
	try {
		for(long npin = 0; pins->Next(1u, &pin, NULL) == S_OK; npin++) {
			try {
				if (SUCCEEDED(pin->QueryDirection(&dir)) && dir == PINDIR_OUTPUT) {
					AM_MEDIA_TYPE *pMediaType = (AM_MEDIA_TYPE *)CoTaskMemAlloc(sizeof(AM_MEDIA_TYPE));
					if (pMediaType != NULL) {
						try {
							if (SUCCEEDED(pin->ConnectionMediaType(pMediaType))) {
								try {
									if ( IsEqualCLSID(pMediaType->majortype, MEDIATYPE_Audio) ||
										 IsEqualCLSID(pMediaType->majortype, MEDIATYPE_AnalogAudio) ||
										 IsEqualCLSID(pMediaType->majortype, MEDIATYPE_Midi) ) {
										PinInfo* info = new PinInfo(npin, pin, pMediaType);
										if (!info)
											res = E_OUTOFMEMORY;
										else {
											pMediaType = NULL;
											this->streams->Add(info);
										}
									}
								} finally (
									if (pMediaType)
										FreePinMediaType(*pMediaType);
								)
							}
						} finally (
							if (pMediaType)
								CoTaskMemFree(pMediaType);
						)
					}
				}
			} finally (
				FreeObject(pin);
			)
		}
	} finally (
		FreeObject(pins);
	)
	_sntprintf(textBuffer, textBufferLength, _T("Found %Id audio streams\r\n"), this->streams->GetCount());
	OutputDebugString(textBuffer);
	return res;
}

HRESULT STDMETHODCALLTYPE AVISplitter::Count(DWORD *pcStreams) {
	if (!pcStreams) return E_POINTER;
	HRESULT res;
	if (FAILED(res = EnsurePins()))
		return res;
	*pcStreams = (DWORD)this->streams->GetCount();
	return S_OK;
}

HRESULT STDMETHODCALLTYPE AVISplitter::Info(long lIndex, AM_MEDIA_TYPE **ppmt, DWORD *pdwFlags, LCID *plcid, DWORD *pdwGroup, LPWSTR *ppszName, IUnknown **ppObject, IUnknown **ppUnk) {
	HRESULT res;
	IEnumPins* pins = NULL;
	IPin* pin = NULL;
	CreateTextBuffer(textBuffer);
	CreateConvertBufferW(convertBuffer);

	if (FAILED(res = EnsurePins()))
		return res;
	if (lIndex < 0 || lIndex >= (long)this->streams->GetCount())
		return S_FALSE;
	PinInfo* info = this->streams->ElementAt(lIndex);
	if (ppmt) {
		AM_MEDIA_TYPE *pMediaType = (AM_MEDIA_TYPE *)CoTaskMemAlloc(sizeof(AM_MEDIA_TYPE));
		if (!pMediaType)
			return E_OUTOFMEMORY;
		if (info->type) {
			if (FAILED(res=CopyPinMediaType(pMediaType, info->type)))
				return res;
		} else {
			ZeroMemory(pMediaType, sizeof(AM_MEDIA_TYPE));
		}
		*ppmt = pMediaType;
	}
	if (pdwFlags) {
		pdwFlags = 0;
	}
	if (plcid) {
		*plcid = 0;
	}
	if (pdwGroup) {
		*pdwGroup = 0;
	}

	if (ppObject) *ppObject = NULL;
	if (ppUnk) *ppUnk = NULL;

	if (ppszName) {
		OutputDebugString(_T("Query name of track"));
		_sntprintf(textBuffer, textBufferLength, _T("Track %ld"), lIndex);
		*ppszName = SysAllocString(ToLPWSTR(textBuffer, convertBuffer, convertBufferLength));
	}

	return S_OK;
}

HRESULT STDMETHODCALLTYPE AVISplitter::Enable(long lIndex, DWORD dwFlags) {
	IPin* connect;
	IMediaControl* control = NULL;
	IMediaSeeking* seek = NULL;
	LONGLONG position = 0;
	HRESULT res;
	LPCTSTR action = NULL;
	CreateTextBuffer(textBuffer);

	switch(dwFlags) {
	case 0:
		return S_OK;
	case AMSTREAMSELECTENABLE_ENABLEALL:
		return E_NOTIMPL;
	}

	if (FAILED(res = EnsurePins()) || FAILED(res = EnsureGraph()))
		return res;
	if (lIndex < 0 || lIndex >= (long)this->streams->GetCount())
		return S_FALSE;

	OAFilterState state = State_Running;
	try {
		if (SUCCEEDED(this->graph->QueryInterface(IID_IMediaControl, (void **)&control))) {
			control->GetState(10, &state);
			if (state != State_Stopped) {
				if (SUCCEEDED(this->graph->QueryInterface(IID_IMediaSeeking, (void **)&seek))) {
					seek->GetCurrentPosition(&position);
				}
				control->Stop();
			}
		}

		for(INT_PTR k = 0; k < this->streams->GetCount(); k++) {
			PinInfo* info = this->streams->ElementAt(k);
			bool enable = dwFlags == AMSTREAMSELECTENABLE_ENABLEALL || 
				dwFlags == AMSTREAMSELECTENABLE_ENABLE && k == lIndex;
			bool connected = false;
			if (SUCCEEDED(info->pin->ConnectedTo(&connect)) && connect) {
				FreeObject(connect);
				connected = true;
			}
			res = S_OK;
			if (enable && !connected) {
				_sntprintf(textBuffer, textBufferLength, _T("Connect pin to %s\r\n"), info->filterName);
				OutputDebugString(textBuffer);
				if (info->pin && info->connect) {
					res = this->graph->ConnectDirect(info->pin, info->connect, info->type);
					action = _T("Connection from splitter to audio decoder");
				}
				if (SUCCEEDED(res) && info->next && info->next->pin && info->next->connect) {
					res = this->graph->ConnectDirect(info->next->pin, info->next->connect, info->next->type);
					action = _T("Connection from audio decoder");
				}
			} else if (!enable && connected) {
				_sntprintf(textBuffer, textBufferLength, _T("Disconnect pin from %s\r\n"), info->filterName);
				OutputDebugString(textBuffer);
				if (info->next && info->next->pin) {
					res = this->graph->Disconnect(info->next->pin);
					action = _T("Disconnection output audio decoder pin");
				}
				if (SUCCEEDED(res) && info->next && info->next->connect) {
					res = this->graph->Disconnect(info->next->connect);
					action = _T("Disconnection input audio pin");
				}
				if (SUCCEEDED(res) && info->pin) {
					res = this->graph->Disconnect(info->pin);
					action = _T("Disconnection output splitter pin");
				}
				if (SUCCEEDED(res) && info->connect) {
					res = this->graph->Disconnect(info->connect);
					action = _T("Disconnection input audio decoder pin");
				}
			}
			if (FAILED(res)) {
				_sntprintf(textBuffer, textBufferLength, _T("%s failed with code %X\r\n"), action, res);
				OutputDebugString(textBuffer);
			}
		}

		if (control && state != State_Stopped) {
			if (seek)
				seek->SetPositions(&position, AM_SEEKING_AbsolutePositioning, NULL, AM_SEEKING_NoPositioning);
			switch(state) {
			case State_Running:
				control->Run();
				break;
			case State_Paused:
				control->Pause();
				break;
			}
		}
	} finally (
		FreeObject(control);
		FreeObject(seek);
	)

	return res;
}

